from mindformers.tools.utils import replace_rank_id_in_ckpt_name
import os
from glob import glob
import shutil

checkpoint_dir = ""
save_prefix = ""


if checkpoint_dir.endswith("/"):
    checkpoint_dir = checkpoint_dir[:-1]
new_checkpoint_dir = checkpoint_dir + "_infer"

rank_dirs = os.listdir(checkpoint_dir)
new_rank_dirs = []
for rank_dir in rank_dirs:
    if rank_dir.startswith("rank_"):
        new_rank_dirs.append(rank_dir)
new_rank_dirs.sort(key=lambda x: int(x.split("rank_")[1]))
rank_dirs = new_rank_dirs
print("rank_dirs: ", rank_dirs)

for rank_dir in rank_dirs:
    print("===========================")
    print(rank_dir)
    rank_id = int(rank_dir.split("rank_")[1])
    prefix = replace_rank_id_in_ckpt_name(save_prefix, rank_id)
    print("prefix: ", prefix)

    new_rank_dir = os.path.join(new_checkpoint_dir, rank_dir)
    os.makedirs(new_rank_dir, exist_ok=True)

    ckpts = glob(os.path.join(checkpoint_dir, rank_dir, "*.ckpt"))
    new_ckpt = None
    for ckpt in ckpts:
        print("ckpt: ", ckpt)
        if os.path.basename(ckpt).startswith(prefix):
            new_ckpt = os.path.join(new_rank_dir, os.path.basename(ckpt))
            break
    if new_ckpt is None:
        raise ValueError("checkpoint not found")
    shutil.copy(ckpt, new_ckpt)
